import os
import torch
import numpy as np
from PIL import Image
import matplotlib.cm as cm

def visual_map(image_path, img, anomaly_map, output_dir, gt_mask=None):
    os.makedirs(output_dir, exist_ok=True)

  
    img_filename = os.path.basename(image_path).split('.')[0]

    # 还原输入图像
    mean = torch.tensor([0.485, 0.456, 0.406], device=img.device).view(3, 1, 1)
    std = torch.tensor([0.229, 0.224, 0.225], device=img.device).view(3, 1, 1)
    img = img * std + mean
    img = (img * 255).clamp(0, 255).byte().cpu().numpy().transpose(1, 2, 0)
    img = Image.fromarray(img)

    # 生成异常图
    anomaly_map = anomaly_map / anomaly_map.max()
    anomaly_map = (cm.jet(anomaly_map)[:, :, :3] * 255).astype(np.uint8)
    anomaly_img = Image.fromarray(anomaly_map)

    # 叠加异常图到原图
    blended_img = Image.blend(img.convert("RGB"), anomaly_img.convert("RGB"), alpha=0.4)

    # 保存结果
    blended_img.save(os.path.join(output_dir, f"{img_filename}.png"))

    # 处理 Ground Truth Mask
    if gt_mask is not None:
        gt_mask_img = (gt_mask * 255).astype(np.uint8)
        gt_mask_img = Image.fromarray(gt_mask_img).convert("RGB")
        gt_mask_img.save(os.path.join(output_dir, f"{img_filename}_gt.png"))
